# --*-- coding: UTF-8 -*-

import sys
sys.path.append('/home/cnn/pro/yazif/vqa/')

from models.vgg_model import VggModel
from utils import config
from utils import image
import tensorflow as tf
import numpy as np
import time
import h5py
import json
import gc

# 该py用来预处理图片，生成训练和测试的图片特征
lstm_config_path = '../configs/lstm.json'
config_json, _ = config.get_config_from_json('../configs/prepro.json')

batch_size = 20



prepro_json,_ = config.get_config_from_json(config_json.json_file_path)

img_path = '%s/%s'

train_img_map = prepro_json.ix_to_img_train
val_img_map = prepro_json.ix_to_img_val



def build_data(img_map, sess, file_path, batch_size=20, name=None):
    sz = len(img_map)
    # sz = 80000
    print ('process %s include %d image, waiting...' % (name, sz))

    f = h5py.File(file_path, 'a')
    f.create_dataset(name=name, shape=(sz, 4096), dtype=np.float32)
    dataset = f[name]

    for left_index in xrange(0, sz, batch_size):

        if left_index % 5000 == 0:
            f.close()
            del dataset
            gc.collect()
            print('complete %d'% (left_index) )
            f = h5py.File(file_path, 'a')
            dataset = f[name]

        # 确定右索引
        right_index = min(sz, left_index + batch_size)
        num = right_index - left_index
        batch = np.zeros([num, 224, 224, 3])
        for index in xrange(0, num):
            batch[index, ...] = image.load_image(img_path % (config_json.img_root, img_map.get(str(left_index + index))))

        feed_dict = {vgg.rgb: batch}

        fc7 = sess.run(vgg.fc7, feed_dict=feed_dict)
        dataset[left_index:right_index, ...] = fc7

        del batch
        del fc7

    if f is not None:
        f.close()
    print ('success!')

with tf.Session(config=tf.ConfigProto(
        gpu_options=(tf.GPUOptions(per_process_gpu_memory_fraction=0.8)))) as sess:
    vgg = VggModel(config_json.vgg19_npy_path)
    vgg.build()

    f = h5py.File(config_json.img_h5_path, 'w')
    f.close()

    start_time = time.time()
    build_data(train_img_map, sess, config_json.img_h5_path, batch_size, 'feat_train')
    build_data(val_img_map, sess, config_json.img_h5_path, batch_size, 'feat_val')

    _, lstm_config = config.get_config_from_json(lstm_config_path)
    lstm_config["is_have_img_data"] = "True"
    lstm_config["img_h5_path"] = config_json.img_h5_path

    json.dump(lstm_config, open(lstm_config_path, 'w'))
    print("wrote ", lstm_config_path)

    print("build finished: %ds" % (time.time() - start_time))

